import numpy as np
import platgo as pg
import scipy.io as sio


class CEC_2020_F3(pg.Problem):

    def __init__(self, D=None) -> None:
        self.name = 'CEC_2020_F3'
        self.type['single'], self.type['real'] = [True] * 2
        self.M = 1
        load_path = 'CEC2020.mat'
        load_data = sio.loadmat(load_path)
        mat = []
        for k in load_data.items():
            mat.append(k)
        self.D = D
        self.O = mat[3][1][0][2][0][0][0]
        if self.D is None or self.D < 10:
            self.D = 5
            self.Mat = mat[3][1][0][2][0][0][1]
        elif self.D < 15:
            self.D = 10
            self.Mat = mat[3][1][0][2][0][0][2]
        elif self.D < 20:
            self.D = 15
            self.Mat = mat[3][1][0][2][0][0][3]
        else:
            self.D = 20
            self.Mat = mat[3][1][0][2][0][0][4]
        lb = [-100] * self.D
        ub = [100] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        s = 1 - 1 / (2 * np.sqrt(self.D + 20) - 8.2)
        mu0 = 2.5
        mu1 = -np.sqrt((mu0 ** 2 - 1) / s)
        Y = (pop.decs - np.tile(self.O[0][0: pop.decs.shape[1]], (pop.decs.shape[0], 1))) / 10
        tmp = 2 * np.tile(np.sign(self.O[0][0: pop.decs.shape[1]]), (pop.decs.shape[0], 1)) * Y + mu0
        Z = np.dot((tmp - mu0), self.Mat.T)
        pop.objv = 700 + np.minimum(np.sum((tmp - mu0) ** 2, axis=1), s * np.sum((tmp - mu1) ** 2, axis=1)) + 10 * (
                self.D - np.sum(np.cos(2 * np.pi * Z), axis=1))
        pop.objv = pop.objv.reshape(pop.objv.shape[0], 1)

    def get_optimal(self) -> np.ndarray:
        pass


if __name__ == '__main__':
    problem = CEC_2020_F3()
    alg = pg.algorithms.GA(problem=problem, maxgen=100)
    pop = alg.go(100)
    print(pop)
